package site.wanjiahao.trie;


import java.util.HashMap;
import java.util.Map;

/**
 * 前缀树，一个节点存放一个字符，可以非常方便判断一个字符是否包含，并且判断前缀
 *
 * @param <V>
 */
public class PrefixTree1<V> implements Trie<V> {

    // 根节点对象
    private Node<V> root;

    private int size;

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public void clear() {
        size = 0;
        // GC Roots至空
        root = null;
    }

    @Override
    public V add(String key, V value) {
        checkPrefixIsLegal(key);
        // 初始化根节点对象
        if (root == null) root = new Node<>(null, null);

        int strLen = key.length();
        Node<V> curNode = root;

        for (int i = 0; i < strLen; i++) {

            char curChar = key.charAt(i);
            Map<Character, Node<V>> children = curNode.children;
            if (children == null) {
                children = new HashMap<>();
                curNode.children = children;
            }

            Node<V> childrenNode = children.get(curChar);
            if (childrenNode == null) {
                childrenNode = new Node<>(curChar, curNode);
                children.put(curChar, childrenNode);
            }

            curNode = childrenNode;
        }
        // 设置最终数据
        curNode.value = value;
        curNode.word = key;
        curNode.isWord = true;
        // 元素自增
        size++;
        return curNode.value;
    }

    @Override
    public boolean contains(String key) {
        Node<V> node = node(key);
        return node != null && node.isWord;
    }

    @Override
    public V get(String key) {
        Node<V> node = node(key);
        if (node != null && node.isWord) return node.value;
        return null;
    }

    @Override
    public V remove(String key) {
        // 获取当前对应的节点
        Node<V> node = node(key);
        V oldValue = node.value;
        if (node != null && node.isWord) {
            // 孩子不为空
            if (node.children != null) {
                size--;
                node.isWord = false;
                return node.value;
            }

            do {
                // 孩子为空, 有节点是单词，或者当前节点还有其他孩子节点
                Node<V> parent = node.parent;
                // 删除当前节点
                parent.children.remove(node.character);
                if (parent.children.size() > 0 || parent.isWord) {
                    size--;
                    return oldValue;
                }
                // 清除当前Map
                parent.children = null;
                node = parent;
            } while (node != null);
        }
        size--;
        return oldValue;
    }

    @Override
    public boolean startsWith(String prefix) {
        return node(prefix) != null;
    }

    private Node<V> node(String key) {
        checkPrefixIsLegal(key);
        // 根据字符串找到对应的节点
        if (root == null) return null;
        int strLen = key.length();

        Node<V> curNode = root;
        for (int i = 0; i < strLen; i++) {
            char curChar = key.charAt(i);
            Map<Character, Node<V>> children = curNode.children;
            if (children == null) return null;
            Node<V> childrenNode = children.get(curChar);
            if (childrenNode == null) return null;
            curNode = childrenNode;
        }

        return curNode;
    }

    private void checkPrefixIsLegal(String prefix) {
        if (prefix == null || prefix.length() == 0)
            throw new IllegalArgumentException("IllegalArgument !");
    }

    // 前缀树节点对象
    private static class Node<V> {

        private V value;

        private String word;

        private boolean isWord;

        private Character character;

        private Node<V> parent;

        private Map<Character, Node<V>> children;


        public Node(Character character, Node<V> parent) {
            this.character = character;
            this.parent = parent;
        }
    }
}
